package dht

import (
	"errors"
	"fmt"
	"log"
	"net"
	"sync"

	"gitee.com/iuhjui/dht/store"
	"github.com/anacrolix/utp"
)

type Handler func(data []byte) ([]byte, error)

type DHT struct {
	conf     *DHTConfig
	socket   *utp.Socket
	router   *Router
	store    store.Store
	sessions map[ID]*Node
	services map[MsgType]Handler
	queue    chan *Node
	mu       sync.Mutex
}

func NewDHT(conf *DHTConfig) (*DHT, error) {
	dht := &DHT{}
	err := conf.Check()
	if err != nil {
		return nil, err
	}
	dht.conf = conf
	err = dht.makeSocket()
	if err != nil {
		return nil, err
	}
	err = dht.makeStore()
	if err != nil {
		return nil, err
	}
	dht.router = NewRouter(conf.ID, 8)
	dht.sessions = map[ID]*Node{}
	dht.services = map[MsgType]Handler{
		PING:         dht.pingHandler,
		FIND_NODE:    dht.findNodeHandler,
		PROVIDE:      dht.provideHandler,
		FIND_PROVIDE: dht.findProvideHandler,
	}
	go dht.serve()
	dht.bootstrap()
	return dht, nil
}

func (dht *DHT) makeSocket() error {
	ip := dht.conf.IP
	port := fmt.Sprint(dht.conf.Port)
	laddr, err := net.ResolveUDPAddr("udp", ip+":"+port)
	if err != nil {
		return err
	}
	conn, err := net.ListenUDP("udp", laddr)
	if err != nil {
		return err
	}
	dht.socket, err = utp.NewSocketFromPacketConn(conn)
	if err != nil {
		return err
	}
	return nil
}

func (dht *DHT) makeStore() error {
	var db store.Store
	if dht.conf.DBType == "memory" {
		db = store.OpenMemoryStore()
	}
	dht.store = db
	return nil
}

func (dht *DHT) serve() {
	for {
		conn, err := dht.socket.Accept()
		if err != nil {
			log.Println("ERROR:", err)
			continue
		}
		go dht.handleQuery(conn)
	}
}

func (dht *DHT) handleQuery(conn net.Conn) {
	defer conn.Close()
	var msg Msg
	req, err := msg.Read(conn)
	if err != nil {
		log.Println("ERROR:", err)
		return
	}
	fn, ok := dht.services[req.Type()]
	if ok != true {
		err := errors.New("service does not exist")
		msg.WriteError(conn, dht.ID(), err)
		log.Println("ERROR:", err)
		return
	}
	buf, err := fn(req.Data())
	if err != nil {
		msg.WriteError(conn, dht.ID(), err)
		log.Println("ERROR:", err)
		return
	}
	node := &Node{
		ID:   req.ID(),
		Addr: conn.RemoteAddr().String(),
	}
	dht.update(node)
	err = msg.Write(conn, dht.ID(), req.Type(), buf)
	if err != nil {
		log.Println("ERROR:", err)
		return
	}
}

func (dht *DHT) query(addr string, mt MsgType, data []byte) ([]byte, error) {
	conn, err := dht.socket.Dial(addr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	var msg Msg
	err = msg.Write(conn, dht.ID(), mt, data)
	if err != nil {
		return nil, err
	}
	resp, err := msg.Read(conn)
	if err != nil {
		return nil, err
	}
	if resp.Type() == ERROR {
		return nil, errors.New(string(resp.Data()))
	}
	return resp.Data(), nil
}

func (dht *DHT) update(node *Node) {
	dht.mu.Lock()
	defer dht.mu.Unlock()
	dht.sessions[node.ID] = node
	dht.router.Update(node.ID)
}

func (dht *DHT) bootstrap() {
	var wg sync.WaitGroup
	for _, addr := range dht.conf.Seeds {
		wg.Add(1)
		go func(addr string) {
			defer wg.Done()
			node, err := dht.Ping(addr)
			if err != nil {
				log.Println("ERROR:", err)
				return
			}
			dht.update(node)
		}(addr)
	}
	wg.Wait()
	list, err := dht.FindNode(dht.Addr().String(), dht.ID())
	if err != nil {
		log.Println("ERROR:", err)
		return
	}
	for _, v := range list {
		dht.update(v)
	}
}

func (dht *DHT) ID() ID {
	return dht.conf.ID
}

func (dht *DHT) Addr() net.Addr {
	return dht.socket.Addr()
}
